{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "96ff9e912bfe9d8",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# viking DB\n",
    "\n",
    ">[viking DB](https://www.volcengine.com/docs/6459/1163946) is a database that stores, indexes, and manages massive embedding vectors generated by deep neural networks and other machine learning (ML) models.\n",
    "\n",
    "This notebook shows how to use functionality related to the VikingDB vector database.\n",
    "\n",
    "To run, you should have a [viking DB instance up and running](https://www.volcengine.com/docs/6459/1165058).\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd771e02d8a93a0",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade volcengine"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12719205caed0d18",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "We want to use VikingDBEmbeddings so we have to get the VikingDB API Key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fbfb32665b4a3640",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-21T09:53:24.186916Z",
     "start_time": "2023-12-21T09:53:24.179524Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OpenAI API Key:\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8c983d329237fa4",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from langchain.document_loaders import TextLoader\n",
    "from langchain.vectorstores.vikingdb import VikingDB, VikingDBConfig\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a4aea2eaeb2261",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "loader = TextLoader(\"./test.txt\")\n",
    "documents = loader.load()\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=10, chunk_overlap=0)\n",
    "docs = text_splitter.split_documents(documents)\n",
    "\n",
    "embeddings = OpenAIEmbeddings()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfd593f3deabfaf8",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "db = VikingDB.from_documents(\n",
    "    docs,\n",
    "    embeddings,\n",
    "    connection_args=VikingDBConfig(\n",
    "        host=\"host\", region=\"region\", ak=\"ak\", sk=\"sk\", scheme=\"http\"\n",
    "    ),\n",
    "    drop_old=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "50e6ee12ca7eec39",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-21T10:01:47.355894Z",
     "start_time": "2023-12-21T10:01:47.334789Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "query = \"What did the president say about Ketanji Brown Jackson\"\n",
    "docs = db.similarity_search(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b6b81f5995c79ef0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-21T10:01:47.771478Z",
     "start_time": "2023-12-21T10:01:47.731485Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "docs[0].page_content"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2d932c1290478ee",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Compartmentalize the data with viking DB Collections\n",
    "\n",
    "You can store different unrelated documents in different collections within same viking DB instance to maintain the context"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "907de4eb10626d2a",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Here's how you can create a new collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f5a59ba40f7985f",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "db = VikingDB.from_documents(\n",
    "    docs,\n",
    "    embeddings,\n",
    "    connection_args=VikingDBConfig(\n",
    "        host=\"host\", region=\"region\", ak=\"ak\", sk=\"sk\", scheme=\"http\"\n",
    "    ),\n",
    "    collection_name=\"collection_1\",\n",
    "    drop_old=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c8eada37b17d992",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "And here is how you retrieve that stored collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "883ec678d47c9adc",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "db = VikingDB.from_documents(\n",
    "    embeddings,\n",
    "    connection_args=VikingDBConfig(\n",
    "        host=\"host\", region=\"region\", ak=\"ak\", sk=\"sk\", scheme=\"http\"\n",
    "    ),\n",
    "    collection_name=\"collection_1\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f0be30cfe70083d",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "After retrieval you can go on querying it as usual."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
